import tensorflow as tf
from utils.ops import get_edit_distance


def train_procedure_keras_model(acoustic_model,data_generator,loss,optimizer):
    data=data_generator.__next__()
    with tf.GradientTape() as tape:
        y_=acoustic_model(data[0])
        l=loss(data[1],y_)
    variables=acoustic_model.trainable_variables
    gradients=tape.gradient(l,variables)
    optimizer.apply_gradients(zip(gradients,variables))
    return l

def cal_train_word_error(accoustic_model,val_data_generator,times,pp_f,pp_t,pp_d):
    if isinstance(accoustic_model,(tf.keras.models.Model)):
        n=0
        costs=0
        for i in range(times):
            data=val_data_generator.__next__()
            y_=accoustic_model(data[0])
            pp_y=pp_f(y_,*pp_t,**pp_d)
            costs+=get_edit_distance(data[1],pp_y)
            for j in range(len(data[1])):
                n+=len(data[1][j])
        return costs/n